# predict_from_trained.py

import pandas as pd
import numpy as np
import os
from datetime import datetime
from sklearn.compose import ColumnTransformer
import joblib
from logUtil import get_logger  # 假设 logUtil.py 在正确的位置

# --- 配置 ---
# 假设此脚本与 train.py 在同一项目根目录下

BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # 指向 train 目录
DATA_DIR = os.path.join(BASE_DIR, '..', 'data')
MODEL_DIR = os.path.join(BASE_DIR, '..')  # 指向 talents 目录
LOG_DIR = os.path.join(BASE_DIR, '..', 'log')  # 指向 talents/log

os.makedirs(LOG_DIR, exist_ok=True)

# 日志
logger = get_logger('predict', log_dir=LOG_DIR)
logger.info(f"Prediction script started at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

# --- 加载模型组件 ---
model_path = os.path.join(MODEL_DIR, 'model/trained_voting_classifier2.pkl')
preprocessor_path = os.path.join(MODEL_DIR, 'model/preprocessor.pkl')
feature_names_path = os.path.join(MODEL_DIR, 'model/feature_names.npy')

# 检查文件是否存在
for path, name in [(model_path, "Model"), (preprocessor_path, "Preprocessor"), (feature_names_path, "Feature Names")]:
    if not os.path.exists(path):
        logger.error(f"{name} file not found: {path}")
        raise FileNotFoundError(f"{name} file not found: {path}")

try:
    # 加载模型
    model = joblib.load(model_path)
    logger.info(f"Model loaded successfully from {model_path}")

    # 加载预处理器
    preprocessor = joblib.load(preprocessor_path)
    logger.info(f"Preprocessor loaded successfully from {preprocessor_path}")

    # 加载原始特征名称
    original_feature_names = np.load(feature_names_path, allow_pickle=True)
    logger.info(f"Original feature names loaded: {original_feature_names}")

except Exception as e:
    logger.error(f"Error loading model components: {e}")
    raise

# --- 加载并准备测试数据 ---
test_data_path = os.path.join(DATA_DIR, 'test.csv') # 直接指向 data/test.csv  # 假设你要预测的数据也在 data/test.csv
if not os.path.exists(test_data_path):
    logger.error(f"Test data file not found: {test_data_path}")
    raise FileNotFoundError(f"Test data file not found: {test_data_path}")

try:
    test_data = pd.read_csv(test_data_path)
    logger.info(f"Test data loaded successfully from {test_data_path}. Shape: {test_data.shape}")
except Exception as e:
    logger.error(f"Error loading test data: {e}")
    raise

# --- 数据预处理 (模仿训练时的步骤) ---
try:
    # 1. 检查是否存在所有需要的特征
    missing_features = set(original_feature_names) - set(test_data.columns)
    if missing_features:
        logger.error(f"Missing features in test data: {missing_features}")
        raise ValueError(f"Missing features in test data: {missing_features}")

    # 2. 选择正确的特征列
    X_test_raw = test_data[original_feature_names].copy()
    logger.info(f"Selected {len(original_feature_names)} features for prediction.")

    # 3. 处理缺失值 (简化处理：用测试集均值填充)
    # 注意：更严谨的做法是保存训练集的均值并在预测时使用
    numeric_columns = X_test_raw.select_dtypes(include=['int64', 'float64']).columns
    for col in numeric_columns:
        if X_test_raw[col].isnull().any():
            fill_value = X_test_raw[col].mean()  # 使用测试集均值（简化）
            X_test_raw[col] = X_test_raw[col].fillna(fill_value)
            logger.warning(
                f"Filled missing values in test column '{col}' with test set mean ({fill_value:.4f}). Consider using training set mean.")

    # 4. 应用训练时的预处理器 (多项式特征 + 标准化)
    # 关键：使用 transform，不是 fit_transform
    X_test_processed = preprocessor.transform(X_test_raw)
    logger.info(f"Test data preprocessed. Shape: {X_test_processed.shape}")

except Exception as e:
    logger.error(f"Error during test data preprocessing: {e}")
    raise

# predict_from_trained.py (续)

# ... (之前的代码保持不变) ...

# --- 进行预测 ---
try:
    logger.info("Starting predictions...")
    # 预测类别
    y_pred = model.predict(X_test_processed)
    # 预测概率 (获取属于类别 '1' 的概率)
    y_proba = model.predict_proba(X_test_processed)[:, 1]

    logger.info("Predictions completed.")
    # --- 输出结果 ---
    print("\n--- Prediction Results ---")
    print(f"Number of samples predicted: {len(y_pred)}")

    # --- 绘制 ROC 曲线 ---
    import matplotlib.pyplot as plt
    from sklearn.metrics import roc_curve, auc

    # 检查测试数据中是否存在真实标签 'Attrition'
    if 'Attrition' in test_data.columns:
        y_true = test_data['Attrition']
        # 确保 y_true 是 numpy 数组且为 int 类型（0, 1）
        y_true = np.array(y_true, dtype=int)

        # 计算 ROC 曲线的点
        fpr, tpr, thresholds = roc_curve(y_true, y_proba)
        # 计算 AUC 值
        roc_auc = auc(fpr, tpr)

        # --- 绘制 ROC 曲线 ---
        plt.figure(figsize=(8, 6))
        # 绘制 ROC 曲线
        plt.plot(fpr, tpr, color='darkorange',
                 lw=2, label='ROC curve (area = %0.4f)' % roc_auc)
        # 绘制对角线 (随机猜测的结果)
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Guess (AUC = 0.5)')
        # 设置图形属性
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate (FPR)')
        plt.ylabel('True Positive Rate (TPR)')
        plt.title('Receiver Operating Characteristic (ROC) Curve')
        plt.legend(loc="lower right")
        # 添加网格
        plt.grid(True)
        # 保存图形到文件 (可选)
        roc_plot_path = os.path.join(MODEL_DIR, 'model/roc_curve.png') # 保存到 model 文件夹
        plt.savefig(roc_plot_path)
        logger.info(f"ROC curve saved to {roc_plot_path}")
        # 显示图形 (在某些环境中可能不需要或无法显示)
        # plt.show()
        plt.close() # 关闭图形以释放内存

        print(f"\nROC AUC Score: {roc_auc:.4f}")
        print("ROC curve has been saved as 'roc_curve.png' in the model directory.")

        # --- 可选：打印分类报告 ---
        from sklearn.metrics import accuracy_score, classification_report
        acc = accuracy_score(y_true, y_pred)
        print(f"\nAccuracy on test set: {acc:.4f}")
        print("\nClassification Report:")
        print(classification_report(y_true, y_pred))

    else:
        print("True labels ('Attrition') not found in test.csv. Cannot plot ROC curve.")
        print("Showing predictions only.")

    # 打印前几个预测结果作为示例
    print("\nFirst 10 Predictions:")
    print("Sample | Predicted Class | Probability of Attrition")
    print("-" * 50)
    for i in range(min(10, len(y_pred))):
        print(f"{i + 1:6} | {y_pred[i]:15} | {y_proba[i]:.4f}")


except Exception as e:
    logger.error(f"Error during prediction or plotting: {e}")
    # 尝试关闭可能打开的 matplotlib 图形
    try:
        plt.close()
    except:
        pass
    raise

logger.info("Prediction script finished.")
